package net.aiion.weave.impl.naming.transformation;

import java.lang.annotation.Annotation;
import java.util.ArrayList;
import java.util.Collection;

import net.aiion.weave.asm.Label;
import net.aiion.weave.asm.Opcodes;
import net.aiion.weave.asm.Type;
import net.aiion.weave.spi.bytecode.Bytecode;
import net.aiion.weave.spi.bytecode.BytecodeContext;
import net.aiion.weave.spi.bytecode.BytecodeException;
import net.aiion.weave.spi.bytecode.BytecodeGenerator;
import net.aiion.weave.spi.bytecode.events.AnnotationStartBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.BytecodeEvent;
import net.aiion.weave.spi.bytecode.events.ClassEndBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.ClassStartBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.MethodEndBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.MethodStartBytecodeEvent;
import net.aiion.weave.spi.bytecode.util.BytecodeTransformation;
import net.aiion.weave.spi.injection.InjectionService;
import net.aiion.weave.spi.naming.Context;
import net.aiion.weave.spi.naming.Contextual;
import net.aiion.weave.spi.naming.annotations.Bound;
import net.aiion.weave.spi.naming.annotations.ConversationScoped;
import net.aiion.weave.spi.naming.annotations.DynamicScoped;
import net.aiion.weave.spi.naming.annotations.GlobalScoped;
import net.aiion.weave.spi.naming.annotations.Named;
import net.aiion.weave.spi.naming.annotations.NotScoped;
import net.aiion.weave.spi.naming.annotations.SessionScoped;
import net.aiion.weave.spi.naming.annotations.Unbound;
import net.aiion.weave.spi.naming.annotations.UserScoped;

//TODO Add inheritance support for any @Bound- and/or @Unbound-annotated methods of the super type.
public class ContextualTransformation extends BytecodeTransformation
{
    private static final String CONTEXT_INTERNAL_NAME = Type.getInternalName(Context.class);

    private static final String CONTEXTUAL_ON_BOUND_DESC = "(L" + CONTEXT_INTERNAL_NAME + ";)V";
    private static final String CONTEXTUAL_ON_UNBOUND_DESC = "(L" + CONTEXT_INTERNAL_NAME + ";)V";

    private Collection<BytecodeEvent> onBoundBytecodeEvents;
    private BytecodeGenerator onBoundBytecodeGenerator;

    private Collection<BytecodeEvent> onUnboundBytecodeEvents;
    private BytecodeGenerator onUnboundBytecodeGenerator;

    private String className;
    private String[] interfaceNames;

    private String contextualClassName;
    private String contextualSimpleClassName;

    private Class<? extends Annotation> scopeAnnotationType;

    private String methodName;
    private String methodDescriptor;

    private boolean inConstructor;

    private boolean hasMethodAnnotations;

    public ContextualTransformation() {
    }

    @Override
    public boolean testRetain(BytecodeContext aBytecodeContext, BytecodeEvent aBytecodeEvent) {
        return aBytecodeEvent.getType() == BytecodeEvent.CLASS_START
            && !aBytecodeEvent.asClassStart().containsInterface(Contextual.class);
    }

    @Override
    public boolean testProcess(BytecodeContext aBytecodeContext, BytecodeEvent aBytecodeEvent) {
        return aBytecodeEvent.getType() == BytecodeEvent.ANNOTATION_START
            && aBytecodeEvent.asAnnotationStart().equalsType(Named.class);
    }

    @Override
    public boolean testSkip(BytecodeContext aBytecodeContext, BytecodeEvent aBytecodeEvent) {
        return aBytecodeEvent.getType() == BytecodeEvent.CLASS_START
            && (aBytecodeEvent.asClassStart().isInterface() || aBytecodeEvent.asClassStart().containsInterface(
                Contextual.class));
    }

    @Override
    public void handleBytecodeEvent(final BytecodeContext aBytecodeContext, final BytecodeEvent aBytecodeEvent) {
        switch (aBytecodeEvent.getType()) {
        case BytecodeEvent.CLASS_START:
            handleClassStartBytecodeEvent(aBytecodeContext, aBytecodeEvent.asClassStart());
            break;
        case BytecodeEvent.ANNOTATION_START:
            handleAnnotationStartBytecodeEvent(aBytecodeContext, aBytecodeEvent.asAnnotationStart());
            break;
        case BytecodeEvent.METHOD_START:
            handleMethodStartBytecodeEvent(aBytecodeContext, aBytecodeEvent.asMethodStart());
            break;
        case BytecodeEvent.METHOD_END:
            handleMethodEndBytecodeEvent(aBytecodeContext, aBytecodeEvent.asMethodEnd());
            break;
        case BytecodeEvent.CLASS_END:
            handleClassEndBytecodeEvent(aBytecodeContext, aBytecodeEvent.asClassEnd());
            break;
        default:
            break;
        }

        aBytecodeContext.forwardBytecodeEvent(aBytecodeEvent);
    }

    private void handleClassStartBytecodeEvent(final BytecodeContext aBytecodeContext,
        final ClassStartBytecodeEvent aBytecodeEvent)
    {
        className = aBytecodeEvent.getName();
        interfaceNames = aBytecodeEvent.getInterfaceNames();
        scopeAnnotationType = DynamicScoped.class;

        contextualClassName = className + "$Contextual";
        contextualSimpleClassName = contextualClassName.substring(contextualClassName.lastIndexOf('/') + 1);

        prepareOnBound();
        prepareOnUnbound();
    }

    private void handleAnnotationStartBytecodeEvent(final BytecodeContext aBytecodeContext,
        final AnnotationStartBytecodeEvent aBytecodeEvent)
    {
        if (aBytecodeEvent.equalsType(NotScoped.class)) {
            scopeAnnotationType = NotScoped.class;
        } else if (aBytecodeEvent.equalsType(DynamicScoped.class)) {
            scopeAnnotationType = DynamicScoped.class;
        } else if (aBytecodeEvent.equalsType(GlobalScoped.class)) {
            scopeAnnotationType = GlobalScoped.class;
        } else if (aBytecodeEvent.equalsType(SessionScoped.class)) {
            scopeAnnotationType = SessionScoped.class;
        } else if (aBytecodeEvent.equalsType(ConversationScoped.class)) {
            scopeAnnotationType = ConversationScoped.class;
        } else if (aBytecodeEvent.equalsType(UserScoped.class)) {
            scopeAnnotationType = UserScoped.class;
        }

        if (methodName == null) {
            return;
        }

        hasMethodAnnotations = true;

        if (!inConstructor && aBytecodeEvent.equalsType(Bound.class)) {
            validateOnBoundMethod();
            extendOnBound();
        }

        if (!inConstructor && aBytecodeEvent.equalsType(Unbound.class)) {
            validateOnUnboundMethod();
            extendOnUnbound();
        }
    }

    private void validateOnBoundMethod() {
        if (inConstructor) {
            throw new BytecodeException(String.format(
                "A @Bound annotated method must not be a constructor in class '%s'.", className.replace('/', '.')));
        }

        if (!CONTEXTUAL_ON_BOUND_DESC.equals(methodDescriptor)) {
            throw new BytecodeException(String.format(
                "A @Bound annotated method must not accept any arguments and must return void in method '%s.%s'.",
                className.replace('/', '.'), methodName));
        }
    }

    private void validateOnUnboundMethod() {
        if (inConstructor) {
            throw new BytecodeException(String.format(
                "A @Unbound annotated method must not be a constructor in class '%s'.", className.replace('/', '.')));
        }

        if (!CONTEXTUAL_ON_UNBOUND_DESC.equals(methodDescriptor)) {
            throw new BytecodeException(String.format(
                "A @Unbound annotated method must not accept any arguments and must return void in method '%s.%s'.",
                className.replace('/', '.'), methodName));
        }
    }

    private void handleMethodStartBytecodeEvent(final BytecodeContext aBytecodeContext,
        final MethodStartBytecodeEvent aBytecodeEvent)
    {
        methodName = aBytecodeEvent.getName();
        methodDescriptor = aBytecodeEvent.getDescriptor();

        inConstructor = aBytecodeEvent.isConstructor();
        if (!inConstructor) {
            return;
        }
    }

    private void handleMethodEndBytecodeEvent(final BytecodeContext aBytecodeContext,
        final MethodEndBytecodeEvent aBytecodeEvent)
    {
        methodName = null;
    }

    private void handleClassEndBytecodeEvent(final BytecodeContext aBytecodeContext,
        final ClassEndBytecodeEvent aBytecodeEvent)
    {
        final BytecodeGenerator bytecodeGenerator = aBytecodeContext.getBytecodeGenerator();
        bytecodeGenerator.referenceInnerClass(contextualClassName, className, contextualSimpleClassName,
            Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC);

        finishOnBound();
        finishOnUnbound();

        final BytecodeGenerator contextualClassBytecodeGenerator = Bytecode.createGenerator(Bytecode.createWriter());

        prepareContextualClass(contextualClassBytecodeGenerator);

        implementGetNames(contextualClassBytecodeGenerator);
        implementGetValue(contextualClassBytecodeGenerator);
        implementGetScope(contextualClassBytecodeGenerator);
        implementOnBound(contextualClassBytecodeGenerator);
        implementOnUnbound(contextualClassBytecodeGenerator);

        finishContextualClass(contextualClassBytecodeGenerator);
    }

    private void prepareContextualClass(final BytecodeGenerator aBytecodeGenerator) {
        aBytecodeGenerator.startClass(Opcodes.ACC_PUBLIC | Opcodes.ACC_SYNTHETIC, contextualClassName, Object.class,
            Contextual.class);

        aBytecodeGenerator.startField(Opcodes.ACC_PRIVATE, "context", Context.class);
        aBytecodeGenerator.endField();

        aBytecodeGenerator.startField(Opcodes.ACC_PRIVATE, "value", String.format("L%s;", className));
        aBytecodeGenerator.endField();

        aBytecodeGenerator.startMethod(Opcodes.ACC_PUBLIC, "<init>", "()V");
        aBytecodeGenerator.startCode();
        aBytecodeGenerator.loadReference(0);
        aBytecodeGenerator.invokeSpecial(Object.class, "<init>", void.class);
        aBytecodeGenerator.returnVoid();
        aBytecodeGenerator.endCode(1, 1);
        aBytecodeGenerator.endMethod();
    }

    private void finishContextualClass(final BytecodeGenerator aBytecodeGenerator) {
        aBytecodeGenerator.endClass();
    }

    private void implementGetNames(final BytecodeGenerator aBytecodeGenerator) {
        aBytecodeGenerator.startMethod(Opcodes.ACC_PUBLIC, "getNames", String[].class);
        aBytecodeGenerator.startCode();

        if (interfaceNames == null) {
            interfaceNames = new String[0];
        }

        aBytecodeGenerator.pushInteger(1 + interfaceNames.length);
        aBytecodeGenerator.newReferenceArray(String.class);

        aBytecodeGenerator.duplicateReference();
        aBytecodeGenerator.pushInteger(0);
        aBytecodeGenerator.pushConstant(className.replace('/', '.'));
        aBytecodeGenerator.storeReferenceArray();

        for (int i = 0; i < interfaceNames.length; i++) {
            aBytecodeGenerator.duplicateReference();
            aBytecodeGenerator.pushInteger(i + 1);
            aBytecodeGenerator.pushConstant(interfaceNames[i].replace('/', '.'));
            aBytecodeGenerator.storeReferenceArray();
        }

        aBytecodeGenerator.returnReference();
        aBytecodeGenerator.endCode(4, 1);
        aBytecodeGenerator.endMethod();
    }

    private void implementGetValue(final BytecodeGenerator aBytecodeGenerator) {
        final Label getValueLabel = new Label();

        aBytecodeGenerator.startMethod(Opcodes.ACC_PUBLIC, "getValue", Object.class);
        aBytecodeGenerator.startCode();

        // if (this.value == null) {
        aBytecodeGenerator.loadReference(0);
        aBytecodeGenerator.loadField(contextualClassName, "value", String.format("L%s;", className));
        aBytecodeGenerator.jumpNotNull(getValueLabel);

        // this.value = ((InjectionService)this.context.lookup(InjectionService.class)).injectInstance(<<className>>);
        aBytecodeGenerator.loadReference(0);
        aBytecodeGenerator.loadReference(0);
        aBytecodeGenerator.loadField(contextualClassName, "context", Context.class);
        aBytecodeGenerator.pushConstant(InjectionService.class);
        aBytecodeGenerator.invokeInterface(Context.class, "lookup", Object.class, Class.class);
        aBytecodeGenerator.castChecked(InjectionService.class);
        aBytecodeGenerator.pushConstant(Type.getType(String.format("L%s;", className)));
        aBytecodeGenerator.invokeInterface(InjectionService.class, "injectInstance", Object.class, Class.class);
        aBytecodeGenerator.storeField(contextualClassName, "value", String.format("L%s;", className));
        // }

        // TODO Implement setting the InterceptionService
        if (hasMethodAnnotations) {
            // We assume, that if any method of the component class is annotated with something, the class will also
            // implement the Interceptable interface.
        }

        aBytecodeGenerator.markLabel(getValueLabel);
        aBytecodeGenerator.markFrameSame();

        // return this.value;
        aBytecodeGenerator.loadReference(0);
        aBytecodeGenerator.loadField(contextualClassName, "value", String.format("L%s;", className));
        aBytecodeGenerator.returnReference();

        aBytecodeGenerator.endCode(3, 1);
        aBytecodeGenerator.endMethod();
    }

    private void implementGetScope(final BytecodeGenerator aBytecodeGenerator) {
        aBytecodeGenerator.startMethod(Opcodes.ACC_PUBLIC, "getScope", Class.class);
        aBytecodeGenerator.startCode();

        aBytecodeGenerator.pushConstant(scopeAnnotationType);
        aBytecodeGenerator.returnReference();

        aBytecodeGenerator.endCode(1, 1);
        aBytecodeGenerator.endMethod();
    }

    private void implementOnBound(final BytecodeGenerator aBytecodeGenerator) {
        aBytecodeGenerator.dispatchBytecodeEvents(onBoundBytecodeEvents);
    }

    private void implementOnUnbound(final BytecodeGenerator aBytecodeGenerator) {
        aBytecodeGenerator.dispatchBytecodeEvents(onUnboundBytecodeEvents);
    }

    private void prepareOnBound() {
        onBoundBytecodeEvents = new ArrayList<>();
        onBoundBytecodeGenerator = Bytecode.createGenerator(onBoundBytecodeEvents);

        onBoundBytecodeGenerator.startMethod(Opcodes.ACC_PUBLIC, "onBound", void.class, Context.class);
        onBoundBytecodeGenerator.startCode();

        onBoundBytecodeGenerator.loadReference(0);
        onBoundBytecodeGenerator.loadReference(1);
        onBoundBytecodeGenerator.storeField(contextualClassName, "context", Context.class);
    }

    private void extendOnBound() {
        onBoundBytecodeGenerator.loadReference(0);
        onBoundBytecodeGenerator.invokeVirtual(contextualClassName, "getValue", Object.class);
        onBoundBytecodeGenerator.invokeVirtual(className, methodName, methodDescriptor);

        if (!methodDescriptor.endsWith(")V")) {
            onBoundBytecodeGenerator.popReference();
        }
    }

    private void finishOnBound() {
        onBoundBytecodeGenerator.returnVoid();
        onBoundBytecodeGenerator.endCode(2, 1);
        onBoundBytecodeGenerator.endMethod();
    }

    private void prepareOnUnbound() {
        onUnboundBytecodeEvents = new ArrayList<>();
        onUnboundBytecodeGenerator = Bytecode.createGenerator(onUnboundBytecodeEvents);

        onUnboundBytecodeGenerator.startMethod(Opcodes.ACC_PUBLIC, "onUnbound", void.class, Context.class);
        onUnboundBytecodeGenerator.startCode();
    }

    private void extendOnUnbound() {
        onUnboundBytecodeGenerator.loadReference(0);
        onUnboundBytecodeGenerator.invokeVirtual(contextualClassName, "getValue", Object.class);
        onUnboundBytecodeGenerator.invokeVirtual(className, methodName, methodDescriptor);
    }

    private void finishOnUnbound() {
        onBoundBytecodeGenerator.loadReference(0);
        onBoundBytecodeGenerator.pushNull();
        onBoundBytecodeGenerator.storeField(contextualClassName, "value", Context.class);

        onBoundBytecodeGenerator.loadReference(0);
        onBoundBytecodeGenerator.pushNull();
        onBoundBytecodeGenerator.storeField(contextualClassName, "context", Context.class);

        onUnboundBytecodeGenerator.returnVoid();
        onUnboundBytecodeGenerator.endCode(2, 1);
        onUnboundBytecodeGenerator.endMethod();
    }
}
